import asyncio
import subprocess
import sys
import traceback
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
import base64

# 检查 Docker 是否可用
def check_docker():
    try:
        result = subprocess.run(["docker", "--version"], 
                               capture_output=True, text=True, timeout=5)
        if result.returncode == 0:
            print(f"Docker 已安装: {result.stdout.strip()}")
            return True
        else:
            print(f"Docker 命令返回错误: {result.stderr}")
            return False
    except Exception as e:
        print(f"检查 Docker 时出错: {e}")
        return False

# 检查 mcp/puppeteer 镜像是否存在
def check_puppeteer_image():
    try:
        result = subprocess.run(["docker", "image", "inspect", "mcp/puppeteer"], 
                               capture_output=True, text=True, timeout=5)
        if result.returncode == 0:
            print("mcp/puppeteer 镜像已存在")
            return True
        else:
            print("mcp/puppeteer 镜像不存在，尝试拉取...")
            pull = subprocess.run(["docker", "pull", "mcp/puppeteer"], 
                                 capture_output=True, text=True, timeout=60)
            if pull.returncode == 0:
                print("成功拉取 mcp/puppeteer 镜像")
                return True
            else:
                print(f"拉取镜像失败: {pull.stderr}")
                return False
    except Exception as e:
        print(f"检查镜像时出错: {e}")
        return False

async def run_search(search_query: str, search_engine="baidu"):
    """
    使用Puppeteer工具执行搜索
    """
    print(f"开始搜索: {search_query}")
    
    # 配置 Puppeteer 服务器参数
    server_params = StdioServerParameters(
        command="docker",
        args=[
            "run", "-i", "--rm", "--init", 
            "--network=host",  # 使用主机网络
            "--dns=8.8.8.8", "--dns=8.8.4.4",  # 使用 Google DNS
            "-e", "DOCKER_CONTAINER=true",
            "--shm-size=1g",  # 增加共享内存大小
            "mcp/puppeteer"
        ],
        env={}  # 可添加额外的环境变量
    )

    try:
        # 连接到 Puppeteer 服务器
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                # 初始化会话
                print("初始化会话...")
                await session.initialize()
                print("会话初始化完成")

                # 列出可用工具（用于确认工具名称）
                tools = await session.list_tools()
                print(f"可用工具: {[tool.name for tool in tools.tools]}")

                # 构建搜索URL
                if search_engine == "baidu":
                    search_url = f"https://www.baidu.com/s?wd={search_query.replace(' ', '+')}"
                elif search_engine == "bing":
                    search_url = f"https://www.bing.com/search?q={search_query.replace(' ', '+')}"
                elif search_engine == "google":
                    search_url = f"https://www.google.com/search?q={search_query.replace(' ', '+')}"
                else:
                    search_url = f"https://www.baidu.com/s?wd={search_query.replace(' ', '+')}"
                
                # 步骤1: 导航到搜索引擎
                print(f"导航到搜索引擎: {search_url}")
                navigate_result = await session.call_tool(
                    "puppeteer_navigate",
                    arguments={
                        "url": search_url,
                        "launchOptions": {
                            "headless": True,
                            "defaultViewport": {"width": 1280, "height": 720}
                        }
                    }
                )
                print(f"导航结果: {navigate_result}")

                # 步骤2: 提取搜索结果 - 根据不同的搜索引擎使用不同的选择器
                print("提取搜索结果...")
                if search_engine == "baidu":
                    evaluate_result = await session.call_tool(
                        "puppeteer_evaluate",
                        arguments={
                            "script": """
                                Array.from(document.querySelectorAll('.result.c-container')).map(container => {
                                    const titleElement = container.querySelector('h3');
                                    const linkElement = container.querySelector('a');
                                    // 直接获取link元素中的href属性
                                    const url = linkElement ? linkElement.getAttribute('href') : '';
                                    
                                    // 获取摘要
                                    const abstractElement = container.querySelector('.c-abstract, .content-right_8Zs40');
                                    const abstract = abstractElement ? abstractElement.textContent.trim() : '';
                                    
                                    return {
                                        title: titleElement ? titleElement.textContent.trim() : '',
                                        link: url,
                                        abstract: abstract
                                    };
                                }).filter(result => result.title && result.link);
                            """
                        }
                    )
                elif search_engine == "bing":
                    # Bing搜索结果选择器
                    evaluate_result = await session.call_tool(
                        "puppeteer_evaluate",
                        arguments={
                            "script": """
                                Array.from(document.querySelectorAll('.b_algo')).map(result => {
                                    const titleElement = result.querySelector('h2 a');
                                    const url = titleElement ? titleElement.href : '';
                                    const abstractElement = result.querySelector('.b_caption p');
                                    
                                    return {
                                        title: titleElement ? titleElement.textContent.trim() : '',
                                        link: url,
                                        abstract: abstractElement ? abstractElement.textContent.trim() : ''
                                    };
                                }).filter(result => result.title && result.link);
                            """
                        }
                    )
                else:
                    # 默认选择器（通用）
                    evaluate_result = await session.call_tool(
                        "puppeteer_evaluate",
                        arguments={
                            "script": """
                                Array.from(document.querySelectorAll('a')).map(link => {
                                    if (link.textContent && link.href && !link.href.includes('javascript:') && link.href.startsWith('http')) {
                                        return {
                                            title: link.textContent.trim(),
                                            link: link.href
                                        };
                                    }
                                    return null;
                                }).filter(Boolean);
                            """
                        }
                    )
                
                # 如果百度搜索结果没有链接，尝试获取data-click属性中的链接
                if search_engine == "baidu" and evaluate_result.content and not evaluate_result.isError:
                    result_data = evaluate_result.content[0].text
                    if "link: \"\"" in result_data or "link: \"\"," in result_data:
                        print("尝试获取百度搜索结果中的data-click属性...")
                        evaluate_result = await session.call_tool(
                            "puppeteer_evaluate",
                            arguments={
                                "script": """
                                    Array.from(document.querySelectorAll('.result.c-container')).map(container => {
                                        const titleElement = container.querySelector('h3');
                                        
                                        // 获取data-click属性，该属性包含URL信息
                                        const dataClick = container.getAttribute('data-click');
                                        let url = '';
                                        
                                        if (dataClick) {
                                            try {
                                                const clickData = JSON.parse(dataClick);
                                                if (clickData.url) {
                                                    url = clickData.url;
                                                }
                                            } catch (e) {
                                                // 解析JSON失败
                                            }
                                        }
                                        
                                        // 如果data-click中没有URL，尝试获取mu属性
                                        if (!url) {
                                            const mu = container.getAttribute('mu');
                                            if (mu) url = mu;
                                        }
                                        
                                        // 如果还没有URL，尝试从链接重定向获取
                                        if (!url) {
                                            const linkElement = container.querySelector('a');
                                            if (linkElement) {
                                                url = linkElement.getAttribute('href') || '';
                                            }
                                        }
                                        
                                        // 获取摘要
                                        const abstractElement = container.querySelector('.c-abstract, .content-right_8Zs40');
                                        const abstract = abstractElement ? abstractElement.textContent.trim() : '';
                                        
                                        return {
                                            title: titleElement ? titleElement.textContent.trim() : '',
                                            link: url,
                                            abstract: abstract
                                        };
                                    }).filter(result => result.title);
                                """
                            }
                        )
                
                print("\n搜索结果:")
                if evaluate_result.content and not evaluate_result.isError:
                    search_results = evaluate_result.content[0].text
                    print(search_results)
                    return search_results
                else:
                    print(f"提取结果失败: {evaluate_result}")
                    
                    # 尝试更简单的脚本
                    print("尝试更简单的脚本...")
                    simple_result = await session.call_tool(
                        "puppeteer_evaluate",
                        arguments={
                            "script": "document.title"
                        }
                    )
                    print(f"页面标题: {simple_result}")
                    
                    # 尝试获取所有链接
                    links_result = await session.call_tool(
                        "puppeteer_evaluate",
                        arguments={
                            "script": """
                                Array.from(document.querySelectorAll('a')).map(a => {
                                    return {
                                        text: a.textContent.trim(),
                                        href: a.href
                                    };
                                }).filter(link => link.text && link.href && link.href.startsWith('http')).slice(0, 20);
                            """
                        }
                    )
                    print(f"页面中的链接: {links_result}")
                    
                    return None
    except Exception as e:
        print(f"搜索过程中出错: {e}")
        print(f"错误详情: {traceback.format_exc()}")
        return None

search_engine_dict={
    "baidu":"https://www.baidu.com",
    "google":"https://www.google.co.jp/",
    "bing":"https://www.bing.com/",
    "yahoo":"https://www.yahoo.co.jp/",
    "yandex":"https://www.yandex.com/",
    "duckduckgo":"https://www.duckduckgo.com/",
}

if __name__ == "__main__":
    print(f"Python 版本: {sys.version}")
    print(f"正在使用 MCP SDK v1.9.1")
    
    try:
        print("\n尝试执行搜索...")
        query = "Python MCP tutorial"
        search_result = asyncio.run(run_search(query, search_engine="baidu"))
        print("\n最终搜索结果:", search_result)
    except Exception as e:
        print(f"程序异常: {e}")
        print(f"异常详情: {traceback.format_exc()}")
        sys.exit(1)